from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.retrieval import create_retrieval_chain
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.runnables import RunnablePassthrough
from langchain_text_splitters import RecursiveCharacterTextSplitter

from ChatGLM_new import tongyi_llm, tongyi_embeddings

prompt = ChatPromptTemplate.from_template("告诉我一个有关 {topic}的笑话")
output_parser = StrOutputParser()

# chain = prompt | tongyi_llm | output_parser
# print(chain.invoke({"topic": "python"}))
# print(chain.input_schema.schema())

print("###############操作单变量输入和输出#############")
vectorstore = Chroma.from_texts(
["小明在华为工作"], embedding=tongyi_embeddings
)
retriever = vectorstore.as_retriever()
template ="""仅根据以下上下文回答问题：
{context}

问题：{question}
"""
prompt = ChatPromptTemplate.from_template(template)

retrieval_chain =(
{"context": retriever,"question": RunnablePassthrough()}
| prompt
| tongyi_llm
| StrOutputParser()
)

# print(retrieval_chain.invoke("小强在哪里工作？"))

print("###########################################################")
loader_txt = TextLoader(r'D:\03study\book\python\python_data_course-main\大模型产品开发导论\云岚宗.txt', encoding='utf8')
docs_txt = loader_txt.load()
text_splitter_txt = RecursiveCharacterTextSplitter(chunk_size=384, chunk_overlap=0,
                                                   separators=["\n\n", "\n", " ", "", "。", "，"])
documents_txt = text_splitter_txt.split_documents(docs_txt)
vectordb = Chroma.from_documents(documents=documents_txt, embedding=tongyi_embeddings)

prompt = ChatPromptTemplate.from_template("""使用下面的语料来回答本模板最末尾的问题。如果你不知道问题的答案，直接回答 "我不知道"，禁止随意编造答案。
        为了保证答案尽可能简洁，你的回答必须不超过三句话，你的回答中不可以带有星号。
        请注意！在每次回答结束之后，你都必须接上 "感谢你的提问" 作为结束语
        以下是一对问题和答案的样例：
            请问：秦始皇的原名是什么
            秦始皇原名嬴政。感谢你的提问。

        以下是语料：
<context>
{context}
</context>

Question: {input}""")
# 创建检索链
document_chain = create_stuff_documents_chain(tongyi_llm, prompt)
#print(document_chain.invoke({"input": "秦始皇的原名是什么", "context": "documents_txt"}))

retriever = vectordb.as_retriever()
retrieval_chain = create_retrieval_chain(retriever, document_chain)
response = retrieval_chain.invoke({
    "input": "纳兰桀是谁?"
})
print(response["answer"])